import random


class Node:
    def __init__(self, data):
        self.data = data
        self.right = None
        self.down = None


class SkipList:
    def __init__(self, max_level):
        list = [i for i in range(max_level)]
        left = [Node(float("-inf")) for _ in range(max_level)]
        right = [Node(float("inf")) for _ in range(max_level)]

        for i in range(max_level - 1):
            left[i].right = right[i]
            left[i].down = left[i + 1]
            right[i].down = right[i + 1]
        left[-1].right = right[-1]
        self.head = left[0]

    def search(self, target: int):
        cur_node = self.head

        while cur_node is not None:
            if cur_node.right.data > target:
                cur_node = cur_node.down
            elif cur_node.right.data < target:
                cur_node = cur_node.right
            else:
                return True
        return False

    def add(self, num: int):
        cur_node = self.head
        stack = []

        while cur_node:
            if cur_node.right.data >= num:
                stack.append(cur_node)
                cur_node = cur_node.down
            else:
                cur_node = cur_node.right
        pre_node = None

        while stack:
            cur_node = stack.pop()
            node = Node(num)
            node.right = cur_node.right
            cur_node.right = node
            if not pre_node:
                pass
            else:
                node.down = pre_node
            pre_node = node

            if random.randint(0, 1):
                break

    def erase(self, num: int):
        cur_node = self.head
        is_removed = False
        while cur_node:
            if cur_node.right.data >= num:
                if cur_node.right.data == num:
                    is_removed = True
                    cur_node.right = cur_node.right.right
                cur_node = cur_node.down
            else:
                cur_node = cur_node.right
        return is_removed


if __name__ == "__main__":
    sp = SkipList(16)

    sp.add(7)
    sp.add(15)
    sp.add(30)
    sp.add(35)
    sp.add(55)
    sp.add(71)
    sp.add(159)

    result = sp.search(22)
    print("search result is:", result)

    result = sp.search(35)
    print("search result is:", result)

    result = sp.erase(151)
    print("erase result is:", result)

    result = sp.erase(159)
    print("erase result is:", result)

    sp.add(31)
    result = sp.search(31)
    print("search result is:", result)
